fen = open('train.tags.zh-en.en', encoding='utf8')
fzh = open('train.tags.zh-en.zh', encoding='utf8')
en_zh = []
while True:
    lz = fzh.readline()
    le = fen.readline()
    if not lz:
        assert not le
        break
    lz, le = lz.strip(), le.strip()
    if lz.startswith('<url>'):
        assert le.startswith('<url>')
        lz = fzh.readline()
        le = fen.readline()
        assert lz.startswith('<keywords>')
        assert le.startswith('<keywords>')
        lz = fzh.readline()
        le = fen.readline()
        assert lz.startswith('<speaker>')
        assert le.startswith('<speaker>')
        lz = fzh.readline()
        le = fen.readline()
        assert lz.startswith('<talkid>')
        assert le.startswith('<talkid>')
        lz = fzh.readline()
        le = fen.readline()
        assert lz.startswith('<title>')
        assert le.startswith('<title>')
        lz = fzh.readline()
        le = fen.readline()
        assert lz.startswith('<description>')
        assert le.startswith('<description>')
    else:
        if not lz:
            assert not le
            break
        lee = []
        for w in le.split(' '):
            w = w.replace('.', '').replace(',', '').lower()
            if w:
                lee.append(w)
        en_zh.append([lee, list(lz)])


print(en_zh[:3])


from tqdm import tqdm
en_words = set()
zh_words = set()
for s in tqdm(en_zh):
    for w in s[0]:
        w = w.replace('.', '').replace(',', '').lower()
        if w:
            en_words.add(w)
    for w in s[1]:
        if w:
            zh_words.add(w)


len(en_words)


en_wl = ['<sos>', '<eos>', '<pad>'] + list(en_words)
zh_wl = ['<sos>', '<eos>', '<pad>'] + list(zh_words)
pad_id = 2


en2id = {}
zh2id = {}
for i, w in enumerate(en_wl):
    en2id[w] = i
for i, w in enumerate(zh_wl):
    zh2id[w] = i


import random
random.shuffle(en_zh)
dl = len(en_zh)
train_set = en_zh[:int(dl*0.8)]
dev_set = en_zh[int(dl*0.8):]


import torch
batch_size = 16
data_workers = 0

class MyDataSet(torch.utils.data.Dataset):
    def __init__(self, examples):
        self.examples = examples

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, index):
        example = self.examples[index]
        s1 = example[0]
        s2 = example[1]
        l1 = len(s1)
        l2 = len(s2)
        return s1, l1, s2, l2, index

def the_collate_fn(batch):
    src = [[0]*batch_size]
    tar = [[0]*batch_size]
    src_max_l = 0
    for b in batch:
        #print(b)
        src_max_l = max(src_max_l, b[1])
    tar_max_l = 0
    for b in batch:
        tar_max_l = max(tar_max_l, b[3])
    for i in range(src_max_l):
        l = []
        for x in batch:
            if i < x[1]:
                #print(x[1])
                l.append(en2id[x[0][i]])
            else:
                l.append(pad_id)
        src.append(l)
    
    for i in range(tar_max_l):
        l = []
        for x in batch:
            if i < x[3]:
                l.append(zh2id[x[2][i]])
            else:
                l.append(pad_id)
        tar.append(l)
    indexs = [b[4] for b in batch]
    src.append([1] * batch_size)
    tar.append([1] * batch_size)
    s1 = torch.LongTensor(src)
    s2 = torch.LongTensor(tar)
    return s1, s2, indexs

train_dataset = MyDataSet(train_set)
train_data_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle = True,
    num_workers=data_workers,
    collate_fn=the_collate_fn,
)



dev_dataset = MyDataSet(dev_set)
dev_data_loader = torch.utils.data.DataLoader(
    dev_dataset,
    batch_size=batch_size,
    shuffle = True,
    num_workers=data_workers,
    collate_fn=the_collate_fn,
)


import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.hid_dim = hid_dim
        self.n_layers = n_layers
        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.rnn = nn.LSTM(emb_dim, hid_dim, n_layers, dropout = dropout)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, src):
        #src = [src len, batch size]
        embedded = self.dropout(self.embedding(src))
        #embedded = [src len, batch size, emb dim]
        outputs, (hidden, cell) = self.rnn(embedded)
        #outputs = [src len, batch size, hid dim * n directions]
        #hidden = [n layers * n directions, batch size, hid dim]
        #cell = [n layers * n directions, batch size, hid dim]
        #outputs are always from the top hidden layer
        return hidden, cell


class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.output_dim = output_dim
        self.hid_dim = hid_dim
        self.n_layers = n_layers
        self.embedding = nn.Embedding(output_dim, emb_dim)
        self.rnn = nn.LSTM(emb_dim, hid_dim, n_layers, dropout = dropout)
        self.fc_out = nn.Linear(hid_dim, output_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, input, hidden, cell):
        
        #input = [batch size]
        #hidden = [n layers * n directions, batch size, hid dim]
        #cell = [n layers * n directions, batch size, hid dim]
        
        #n directions in the decoder will both always be 1, therefore:
        #hidden = [n layers, batch size, hid dim]
        #context = [n layers, batch size, hid dim]
        
        input = input.unsqueeze(0)
        
        #input = [1, batch size]
        
        embedded = self.dropout(self.embedding(input))
        
        #embedded = [1, batch size, emb dim]
                
        output, (hidden, cell) = self.rnn(embedded, (hidden, cell))
        
        #output = [seq len, batch size, hid dim * n directions]
        #hidden = [n layers * n directions, batch size, hid dim]
        #cell = [n layers * n directions, batch size, hid dim]
        
        #seq len and n directions will always be 1 in the decoder, therefore:
        #output = [1, batch size, hid dim]
        #hidden = [n layers, batch size, hid dim]
        #cell = [n layers, batch size, hid dim]
        
        prediction = self.fc_out(output.squeeze(0))
        
        #prediction = [batch size, output dim]
        
        return prediction, hidden, cell


class Seq2Seq(nn.Module):
    def __init__(self, input_word_count, output_word_count, encode_dim, decode_dim, hidden_dim, n_layers, encode_dropout, decode_dropout, device):
        super().__init__()
        self.encoder = Encoder(input_word_count, encode_dim, hidden_dim, n_layers, encode_dropout)
        self.decoder = Decoder(output_word_count, decode_dim, hidden_dim, n_layers, decode_dropout)
        self.device = device
        
    def forward(self, src, trg, teacher_forcing_ratio = 0.5):
        #src = [src len, batch size]
        #trg = [trg len, batch size]
        #teacher_forcing_ratio is probability to use teacher forcing
        #e.g. if teacher_forcing_ratio is 0.75 we use ground-truth inputs 75% of the time
        batch_size = trg.shape[1]
        trg_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim
        #tensor to store decoder outputs
        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)

        #last hidden state of the encoder is used as the initial hidden state of the decoder
        hidden, cell = self.encoder(src)

        #first input to the decoder is the <sos> tokens
        input = trg[0,:]

        for t in range(1, trg_len):

            #insert input token embedding, previous hidden and previous cell states
            #receive output tensor (predictions) and new hidden and cell states
            output, hidden, cell = self.decoder(input, hidden, cell)

            #place predictions in a tensor holding predictions for each token
            outputs[t] = output

            #decide if we are going to use teacher forcing or not
            teacher_force = random.random() < teacher_forcing_ratio

            #get the highest predicted token from our predictions
            top1 = output.argmax(1) 

            #if teacher forcing, use actual next token as next input
            #if not, use predicted token
            input = trg[t] if teacher_force else top1
        
        return outputs


source_word_count = len(en_wl)
target_word_count = len(zh_wl)
encode_dim = 256
decode_dim = 256
hidden_dim = 512
n_layers = 2
encode_dropout = 0.5
decode_dropout = 0.5
device = torch.device('cpu')
model = Seq2Seq(source_word_count, target_word_count, encode_dim, decode_dim, hidden_dim, n_layers, encode_dropout, decode_dropout, device).to(device)


def init_weights(m):
    for name, param in m.named_parameters():
        nn.init.uniform_(param.data, -0.08, 0.08)
        
model.apply(init_weights)


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f'The model has {count_parameters(model):,} trainable parameters')


import torch.optim as optim
optimizer = optim.Adam(model.parameters())


criterion = nn.CrossEntropyLoss(ignore_index = pad_id)


def train(model, iterator, optimizer, criterion, clip):
    model.train()
    epoch_loss = 0
    
    for i, batch in enumerate(iterator):
        
        src = batch[0]
        trg = batch[1]
        
        optimizer.zero_grad()
        
        output = model(src, trg)
        
        #trg = [trg len, batch size]
        #output = [trg len, batch size, output dim]
        
        output_dim = output.shape[-1]
        
        output = output[1:].view(-1, output_dim)
        trg = trg[1:].view(-1)
        
        #trg = [(trg len - 1) * batch size]
        #output = [(trg len - 1) * batch size, output dim]
        
        loss = criterion(output, trg)
        
        loss.backward()
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        
        optimizer.step()
        
        epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)


def evaluate(model, iterator, criterion):
    
    model.eval()
    
    epoch_loss = 0
    
    with torch.no_grad():
    
        for i, batch in enumerate(iterator):

            src = batch.src
            trg = batch.trg

            output = model(src, trg, 0) #turn off teacher forcing

            #trg = [trg len, batch size]
            #output = [trg len, batch size, output dim]

            output_dim = output.shape[-1]
            
            output = output[1:].view(-1, output_dim)
            trg = trg[1:].view(-1)

            #trg = [(trg len - 1) * batch size]
            #output = [(trg len - 1) * batch size, output dim]

            loss = criterion(output, trg)
            
            epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)


def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs


import math
import time
N_EPOCHS = 10
CLIP = 1

best_valid_loss = float('inf')

for epoch in range(N_EPOCHS):
    
    start_time = time.time()
    
    train_loss = train(model, train_data_loader, optimizer, criterion, CLIP)
    valid_loss = evaluate(model, dev_data_loader, criterion)
    
    end_time = time.time()
    
    epoch_mins, epoch_secs = epoch_time(start_time, end_time)
    
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'tut1-model.pt')
    
    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')
    print(f'\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')
    print(f'\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')



